package com.xiyuan.smartutils.crypto;

import com.xiyuan.smartutils.*;
import com.xiyuan.smartutils.constants.HttpConstants;
import com.xiyuan.smartutils.reflect.Unsafe;

import javax.net.ssl.*;
import javax.net.ssl.SSLEngineResult.HandshakeStatus;
import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.net.SocketAddress;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.file.Paths;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.security.cert.CRL;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

/**
 * 1)SSL相关工具类<br>
 * 2)SSL参数配置和管理器<br>
 *
 * @version v1.0.0 @author lgz 2017-8-27 新建与整理
 */
@SuppressWarnings({"BooleanMethodIsAlwaysInverted", "AlibabaLowerCamelCaseVariableNaming"})
public class SSL implements HttpConstants {
    static {
        Unsafe.init();
    }
    
    /**
     * 协议类型 默认SSl协议协议
     */
    private String sslProtocol = "SSL";
    
    /**********************************************************/
    /********** 两项必须（服务端公私钥） ****************************/
    /**********************************************************/
    /**
     * keystoreFile 存储位置 例如：./conf/keystore.jks
     */
    private String keystoreFile;
    
    /**
     * keystoreFile 存储内容
     */
    private String keystoreContent;
    
    /**
     * keystoreKey 存储密码 例如：123456
     */
    private String keystoreKey = "";
    
    /********************************************************/
    /********** 三项根据要求可选（CA信任公钥） **********************/
    /********************************************************/
    
    /**
     * truststoreFile 存储位置 例如：./conf/truststore.jks
     */
    private String truststoreFile;
    /**
     * + truststoreFile 存储内容
     */
    private String truststoreContent;
    /**
     * truststoreKey 存储密码 例如：123456
     */
    private String truststoreKey = ""; // = "123456";
    
    /**
     * 是否需要客户端证验证（双向验证，客户端也要导入证书才行）
     */
    private boolean clientAuth; // 是否需要客户端证验证（双向验证，客户端也要导入证书才行）
    
    private SSLContext context;
    
    public SSL() {
    }
    
    /***
     * 初始化参数<br>
     * 参数表（paramMap）包含以下参数：<br>
     * 1）keystoreFile 服务端公私钥文件路径<br>
     * 2）keystoreKey 服务端公私钥<br>
     * 3）truststoreFile CA信任公钥文件路径<br>
     * 4）truststoreKey CA信任公钥<br>
     * 5）clientAuth 是否需要客户端证验证（双向验证，客户端也要导入证书才行）<br>
     *
     * @param paramMap 参数表
     * @throws Exception 异常
     */
    
    public SSL(Map<String, String> paramMap) throws Exception {
        if (paramMap == null) {
            paramMap = new HashMap<>();
        }
        // 1.SSL参数（五项，其中Keystore两项必须）
        this.keystoreFile = paramMap.get(_KEY_STORE_FILE_);
        this.keystoreKey = paramMap.get(_KEY_STORE_KEY_);
        this.sslProtocol = paramMap.get(_SSL_PROTOCOL_);
        
        Asserts.as(Validates.isNotEmpty(keystoreFile) ? null : "配置HTTPS时要求必须配置[keystoreFile]不能为空");
        Asserts.as(Validates.isNotEmpty(keystoreKey) ? null : "配置HTTPS时要求必须配置[keystoreKey]不能为空");
        
        this.truststoreFile = paramMap.get(_TRUST_STORE_FILE_);
        this.truststoreKey = paramMap.get(_TRUST_STORE_KEY_);
        
        Asserts.as((!Validates.isNotEmpty(truststoreFile) || Validates.isNotEmpty(truststoreKey))
                   ? null
                   : "配置HTTPS时如要配置了[truststoreFile]则[truststoreKey]不能为空");
        
        String ret = paramMap.get(_CLIENT_AUTH_);
        
        this.clientAuth = "1".equals(ret) || "true".equalsIgnoreCase(ret);
        this.context = createSSLContext();
    }
    
    /**
     * @param keystoreFile keystoreFile 服务端公私钥文件路径
     * @param keystoreKey  keystoreKey 服务端公私钥
     * @return SSL
     */
    public SSL setKeystoreFile(String keystoreFile, String keystoreKey) throws IllegalArgumentException {
        if (Validates.isBlank((keystoreFile))) {
            throw new IllegalArgumentException("配置HTTPS时要求必须配置[keystoreFile]不能为空");
        }
        File file = new File(keystoreFile);
        if (!file.exists() || !file.isFile() || !file.canRead()) {
            throw new IllegalArgumentException("HTTPS配置[keystoreFile]文件不存在，或者不可读");
        }
        
        return this.setKeystoreContent(Files.readUTF8(file), keystoreKey);
    }
    
    /**
     * @param truststoreFile CA信任公钥，证书文件路径
     * @param truststoreKey  CA信任公钥
     * @return SSL
     */
    public SSL setTruststoreFile(String truststoreFile, String truststoreKey) throws IllegalArgumentException {
        if (Validates.isBlank((truststoreFile))) {
            throw new IllegalArgumentException("配置HTTPS时要求必须配置[truststoreFile]不能为空");
        }
        File file = new File(truststoreFile);
        if (!file.exists() || !file.isFile() || !file.canRead()) {
            throw new IllegalArgumentException("HTTPS配置[truststoreFile]文件不存在，或者不可读");
        }
        
        return this.setTruststoreContent(Files.readUTF8(file), truststoreKey);
    }
    
    public SSL setKeystoreContent(String keystoreContent, String keystoreKey) throws IllegalArgumentException {
        if (Validates.isBlank((keystoreContent))) {
            throw new IllegalArgumentException("配置HTTPS时要求必须配置[keystoreContent]不能为空");
        }
        
        this.keystoreContent = keystoreContent;
        this.keystoreKey = StrUtils.trim(keystoreKey, "");
        return this;
    }
    
    public SSL setTruststoreContent(String truststoreContent, String truststoreKey) throws IllegalArgumentException {
        if (Validates.isBlank((truststoreContent))) {
            throw new IllegalArgumentException("配置HTTPS时要求必须配置[truststoreContent]不能为空");
        }
        
        this.truststoreContent = truststoreContent;
        this.truststoreKey = StrUtils.trim(truststoreKey, "");
        return this;
    }
    
    /**
     * @param clientAuth 是否需要客户端证验证（双向验证，客户端也要导入证书才行）
     * @return SSL
     */
    public SSL setClientAuth(boolean clientAuth) {
        this.clientAuth = clientAuth;
        return this;
    }
    
    /**
     * @param protocol 协议
     * @return SSL
     */
    public SSL setSSLProtocol(String protocol) {
        this.sslProtocol = protocol;
        return this;
    }
    
    /******************************************************************************/
    // 获取SSL创建的SSLContext&SSLServerSocketFactory
    
    /******************************************************************************/
    
    public String getKeystoreFile() {
        return keystoreFile;
    }
    
    public String getKeystoreKey() {
        return keystoreKey;
    }
    
    public String getTruststoreFile() {
        return truststoreFile;
    }
    
    public String getTruststoreKey() {
        return truststoreKey;
    }
    
    public boolean isClientAuth() {
        return clientAuth;
    }
    
    public void setContext(SSLContext context) {
        this.context = context;
    }
    
    public String getSslProtocol() {
        return sslProtocol;
    }
    
    private boolean check() {
        if (Validates.isBlank(keystoreContent)) {
            return false;
        }
        
        return this.truststoreContent != null || this.truststoreKey.length() <= 0;
    }
    
    /**
     * 创建SSL工厂
     */
    public SSLServerSocketFactory createFactory() throws Exception {
        if (!check()) {
            // 调用方法前必须初始化参数
            throw new IllegalStateException("Parameters must be initialized before calling a method");
        }
        
        return getContext().getServerSocketFactory();
    }
    
    public SSLContext getContext() throws Exception {
        return context != null ? context : createSSLContext();
    }
    
    /**
     * 创建SSL环境
     */
    public SSLContext createSSLContext() throws Exception {
        if (!check()) {
            // 调用方法前必须初始化参数
            throw new IllegalStateException("Parameters must be initialized before calling a method");
        }
        
        // .jks
        KeyStore keyStore = KeyStore.getInstance(SSL_KEYSTORE_TYPE);
        keyStore.load(new ByteArrayInputStream(keystoreContent.getBytes()), keystoreKey.toCharArray());
        
        KeyManagerFactory keyManagerFactory = KeyManagerFactory.getInstance(SSL_ALGORITHM);
        keyManagerFactory.init(keyStore, keystoreKey.toCharArray());
        
        TrustManager[] tms = null;
        if (Validates.isNotBlank(truststoreContent)) {
            // 有CA信任证书的创建管理器
            
            // .jks
            KeyStore trustStore = KeyStore.getInstance(SSL_KEYSTORE_TYPE);
            trustStore.load(new ByteArrayInputStream(truststoreContent.getBytes()), truststoreKey.toCharArray());
            
            TrustManagerFactory trustManagerFactory = TrustManagerFactory.getInstance(SSL_ALGORITHM);
            trustManagerFactory.init(trustStore);
            tms = trustManagerFactory.getTrustManagers();
        }
        
        SSLContext sslContext = SSLContext.getInstance(Validates.isBlank(this.sslProtocol)
                                                       ? SSL_PROTOCOL
                                                       : this.sslProtocol);
        
        sslContext.init(keyManagerFactory.getKeyManagers(), tms, new SecureRandom());
        return sslContext;
    }
    
    /******************************************************************************/
    // SSL握手
    /******************************************************************************/
    
    /**
     * AIO 协议握手
     *
     * @return ByteBuffer 握手成功后的解包数据
     */
    public ByteBuffer doHandshake(AsynchronousSocketChannel socket, SSLEngine sslEngine) throws Exception {
        return executeHandshake(new NetSocket(socket), sslEngine, 10);
    }
    
    /**
     * NIO 协议握手
     *
     * @param socket    客户端
     * @param sslEngine SSLEngine
     * @return ByteBuffer 握手成功后的解包数据
     * @throws Exception 异常
     */
    public ByteBuffer doHandshake(SocketChannel socket, SSLEngine sslEngine) throws Exception {
        return executeHandshake(new NetSocket(socket), sslEngine, 10);
    }
    
    
    /**
     * NIO 协议握手
     *
     * @param socket    客户端
     * @param sslEngine SSLEngine
     * @param timeout   握手时长单位/秒，至少4秒
     * @return ByteBuffer 握手成功后的解包数据
     * @throws Exception 异常
     */
    public ByteBuffer doHandshake(SocketChannel socket, SSLEngine sslEngine, int timeout) throws Exception {
        return executeHandshake(new NetSocket(socket), sslEngine, timeout);
    }
    
    /**
     * AIO 协议握手
     *
     * @param socket    客户端
     * @param sslEngine SSLEngine
     * @param timeout   握手时长单位/秒，至少4秒
     * @return ByteBuffer 握手成功后的解包数据
     * @throws Exception 异常
     */
    public ByteBuffer doHandshake(AsynchronousSocketChannel socket, SSLEngine sslEngine, int timeout) throws Exception {
        return executeHandshake(new NetSocket(socket), sslEngine, timeout);
    }
    
    /**
     * SSL 协议握手
     *
     * @param socket    客户端
     * @param sslEngine SSLEngine
     * @param timeout   握手时长单位/秒，至少4秒
     * @return ByteBuffer 第一字节表示握手状态，0表示握手失败，1表示https握手成功，2表示该握手成功，该链接为http链接，后续字节表示应用数据,
     * @throws Exception 异常
     */
    private ByteBuffer executeHandshake(NetSocket socket, SSLEngine sslEngine, int timeout) throws
                                                                                            Exception {
        timeout = Math.max(timeout, 4);
        // 1.初始化参数
        SSLSession session = sslEngine.getSession();
        int appSize = session.getApplicationBufferSize();
        int packetSize = session.getPacketBufferSize();
        
        ByteBuffer appBuff = ByteBuffer.allocate(appSize + 10);
        ByteBuffer packet = ByteBuffer.allocate(packetSize);
        
        // 2.开启握手
        sslEngine.beginHandshake();
        
        // 3 执行握手
        HandshakeStatus hsStatus = sslEngine.getHandshakeStatus();
        boolean handshakeDone = false;
        long tim = System.currentTimeMillis();
        int count = 0;
        boolean isHttp = true;
        String[] http_methods = new String[]{"GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "TRACE"};
        while (!handshakeDone) {
            switch (hsStatus) {
                case NEED_UNWRAP: {
                    // 1. 读数据 客户端请求Hello
                    int len = socket.read((ByteBuffer) packet.clear());
                    
                    // 缓冲切换为读模式
                    packet.flip();
                    if (len < 0) {
                        throw new SocketException("socket is Closed");
                    }
                    
                    // if (isHttp && allowHttp && len > 0) {
                    //     // 判断是否http请求， 标记下次读取则不是http
                    //     isHttp = false;
                    //
                    //     String src = new String(packet.array(), 0, len);
                    //     int ind = src.indexOf("\r\n");
                    //     if (ind != -1) {
                    //         src = src.substring(0, ind);
                    //     }
                    //
                    //     String[] arr = src.split(" ");
                    //     if (arr.length > 2) {
                    //         // [method][uri][HTTP/version]
                    //
                    //         String method = arr[0].toUpperCase();
                    //         String uri = arr[1];
                    //         String version = arr[2];
                    //         boolean flag2 = StrUtils.startsWithIgnoreCase(version, _HTTP_);
                    //         if (!flag2) {
                    //             flag2 = StrUtils.startsWithIgnoreCase(version, _WS_);
                    //         }
                    //
                    //         if (!flag2) {
                    //             // 再次判断 请求方法，如果url 过长 那么第一个packet 装不下，那么 最后一个分片肯定不是HTTP
                    //             for (int i = 0; i < arr.length; i++) {
                    //                 String m = arr[i];
                    //                 if (m.equals(method)) {
                    //                     flag2 = uri != null && uri.charAt(0) == '/';
                    //                     break;
                    //                 }
                    //             }
                    //         }
                    //
                    //         if (flag2) {
                    //             // http 链接 直接返回
                    //             return (ByteBuffer) ByteBuffer.allocate(len).put(packet).flip();
                    //         }
                    //     }
                    // }
                    
                    do {
                        sslEngine.unwrap(packet, (ByteBuffer) appBuff.clear());
                        hsStatus = doTask(sslEngine);
                    }
                    while (packet.hasRemaining() && hsStatus == HandshakeStatus.NEED_UNWRAP);
                    
                    // 判断是否超时
                    if (len == 0 && hsStatus == HandshakeStatus.NEED_UNWRAP) {
                        Threads.sleepIgnoreException(50);
                        long tim2 = System.currentTimeMillis();
                        if (tim2 - tim >= 999) {
                            // 1有一秒钟了
                            count++;
                            tim = tim2;
                        }
                        
                        if (count > timeout) {
                            String str = StrUtils.format("Handshake timeout by %s second", timeout);
                            throw new SocketTimeoutException(str);
                        }
                    }
                    break;
                }
                case NEED_WRAP: {
                    // 2.服务端响应Hello
                    // 写入缓冲
                    sslEngine.wrap(_HELLO_N_, (ByteBuffer) packet.clear());
                    // 响应ssl 数据
                    socket.write((ByteBuffer) packet.flip());
                    hsStatus = doTask(sslEngine);
                    packet.clear().flip();
                    break;
                }
                case NEED_TASK: {
                    // 3.等待任务，当前使用同步处理
                    hsStatus = doTask(sslEngine);
                    break;
                }
                case FINISHED: {
                    // 4.握手完成
                    break;
                }
                case NOT_HANDSHAKING: {
                    // 5.握手成功
                    handshakeDone = true;
                    break;
                }
            }
        }
        
        
        return packet;
    }
    
    /**
     * 同步获取握手状态
     */
    private HandshakeStatus doTask(SSLEngine sslEngine) {
        Runnable runnable;
        while ((runnable = sslEngine.getDelegatedTask()) != null) {
            runnable.run();
        }
        
        return sslEngine.getHandshakeStatus();
    }
    
    /****************************************************/
    // 一下是提供静态调用的函数
    /****************************************************/
    /**
     * 创建默认的客户端SSLSocket<br>
     * java 1.7及以上的版本使用 TLSv1.2协议 以下的使用 TLSv1协议
     */
    public static SSLSocket getSocket(TrustManager trustManager) throws
                                                                 NoSuchAlgorithmException,
                                                                 KeyManagementException,
                                                                 IOException {
        boolean overJava7 = Systems.isJavaVersionMoreThen7();
        return (SSLSocket) getSocketFactory(trustManager, overJava7 ? "TLSv1.2" : "TLSv1").createSocket();
    }
    
    /**
     * 初始化SSL上下文环境
     *
     * @param trustManager
     * @param provider     连接协议，如，TLSv1.2、TLSv1
     * @return SSLSocketFactory
     * @throws NoSuchAlgorithmException 算法异常
     * @throws KeyManagementException   异常
     */
    public static SSLSocketFactory getSocketFactory(TrustManager trustManager, String provider) throws
                                                                                                NoSuchAlgorithmException,
                                                                                                KeyManagementException {
        // 为支持JDK1.6，设置协议为TLSv1，JDK1.7开始有TLSv1.1,TLSv1.2
        SSLContext ctx = SSLContext.getInstance(provider);
        ctx.init(null,
                 new TrustManager[]{trustManager != null ? trustManager : new DefaultTrustManager()},
                 new SecureRandom());
        return ctx.getSocketFactory();
    }
    
    /**
     * 缺省的信任管理器
     */
    public static class DefaultTrustManager implements X509TrustManager {
        public X509Certificate[] getAcceptedIssuers() {
            return null;
        }
        
        public void checkClientTrusted(X509Certificate[] chain, String authType) throws CertificateException {
            // System.out.println("checkClientTrusted");
        }
        
        public void checkServerTrusted(X509Certificate[] chain, String authType) throws CertificateException {
            // System.out.println("checkServerTrusted");
        }
    }
    
    /**
     * 缺省的主机验证类
     */
    public static class DefaultHostnameVerifier implements HostnameVerifier {
        public boolean verify(String hostname, SSLSession session) {
            return true;
        }
    }
    
    /**
     * 推断密码长度SSL
     *
     * @param cipherSuite 原值
     * @return 长度
     */
    public static int deduceKeyLength(String cipherSuite) {
        // 定义最常见和最不常见的
        if (cipherSuite == null) {
            return 0;
        }
        else if (cipherSuite.contains("WITH_RC4_128_")) {
            return 128;
        }
        else if (cipherSuite.contains("WITH_AES_128_")) {
            return 128;
        }
        else if (cipherSuite.contains("WITH_RC4_40_")) {
            return 40;
        }
        else if (cipherSuite.contains("WITH_3DES_EDE_CBC_")) {
            return 168;
        }
        else if (cipherSuite.contains("WITH_IDEA_CBC_")) {
            return 128;
        }
        else if (cipherSuite.contains("WITH_RC2_CBC_40_")) {
            return 40;
        }
        else if (cipherSuite.contains("WITH_DES40_CBC_")) {
            return 40;
        }
        else if (cipherSuite.contains("WITH_DES_CBC_")) {
            return 56;
        }
        else {
            return 0;
        }
    }
    
    /**
     * 通过SSLSession获取X509证书链
     *
     * @param sslSession SSL会话
     * @return X509证书链
     */
    public static X509Certificate[] getCertChain(SSLSession sslSession) {
        try {
            javax.security.cert.X509Certificate[] javaxCerts = sslSession.getPeerCertificateChain();
            if (javaxCerts == null || javaxCerts.length == 0) {
                return null;
            }
            
            int length = javaxCerts.length;
            X509Certificate[] javaCerts = new X509Certificate[length];
            
            CertificateFactory cf = CertificateFactory.getInstance("X.509");
            for (int i = 0; i < length; i++) {
                byte[] bytes = javaxCerts[i].getEncoded();
                ByteArrayInputStream stream = new ByteArrayInputStream(bytes);
                javaCerts[i] = (X509Certificate) cf.generateCertificate(stream);
            }
            
            return javaCerts;
        }
        catch (Exception pue) {
            // pue.printStackTrace();
            return null;
        }
    }
    
    /**
     * 获取KeyStore，指定存储路径、类型、提供者和密码
     *
     * @param storePath     存储路径
     * @param storeType     存储类型
     * @param storeProvider 提供者
     * @param storePassword 密码
     * @return KeyStore
     * @throws Exception 捕捉异常
     */
    public static KeyStore getKeyStore(String storePath, String storeType, String storeProvider,
                                       String storePassword) throws Exception {
        try (InputStream is = java.nio.file.Files.newInputStream(Paths.get(storePath))) {
            return getKeyStore(is, storeType, storeProvider, storePassword);
        }
    }
    
    /**
     * 获取KeyStore，指定流，存储类型、提供者和密码
     *
     * @param is
     * @param storeType     存储类型
     * @param storeProvider 提供者
     * @param storePassword 密码
     * @return KeyStore
     * @throws Exception 捕捉异常
     */
    public static KeyStore getKeyStore(InputStream is, String storeType, String storeProvider,
                                       String storePassword) throws Exception {
        KeyStore keystore = null;
        if (storeProvider != null) {
            keystore = KeyStore.getInstance(storeType, storeProvider);
        }
        else {
            keystore = KeyStore.getInstance(storeType);
        }
        
        keystore.load(is, storePassword == null ? null : storePassword.toCharArray());
        return keystore;
        
    }
    
    /**
     * 加载CRL，指定CRL路径
     *
     * @param crlPath CRL路径
     * @return 得到CRL列表
     * @throws Exception 捕捉异常
     */
    public static Collection<? extends CRL> loadCRL(String crlPath) throws Exception {
        if (crlPath == null) {
            return null;
        }
        
        try (InputStream is = java.nio.file.Files.newInputStream(Paths.get(crlPath))) {
            return CertificateFactory.getInstance("X.509").generateCRLs(is);
        }
    }
    
    static class NetSocket {
        private final SocketChannel nio;
        private final AsynchronousSocketChannel aio;
        
        private final boolean isNio;
        
        private NetSocket(SocketChannel socket) {
            this.isNio = true;
            this.nio = socket;
            this.aio = null;
        }
        
        private NetSocket(AsynchronousSocketChannel socket) {
            this.isNio = false;
            this.nio = null;
            this.aio = socket;
        }
        
        public static NetSocket from(SocketChannel socket) {
            return new NetSocket(socket);
        }
        
        public static NetSocket from(AsynchronousSocketChannel socket) {
            return new NetSocket(socket);
        }
        
        public int write(ByteBuffer buffer) throws Exception {
            if (isNio) {
                return nio.write(buffer);
            }
            else {
                return aio.write(buffer).get();
            }
        }
        
        public int read(ByteBuffer buffer) throws Exception {
            if (isNio) {
                return nio.read(buffer);
            }
            else {
                return aio.read(buffer).get();
            }
        }
        
        public boolean isOpen() {
            if (isNio) {
                return nio.isOpen();
            }
            else {
                return aio.isOpen();
            }
        }
        
        public NetSocket bind(SocketAddress local)
                throws IOException {
            if (isNio) {
                nio.bind(local);
            }
            else {
                aio.bind(local);
            }
            return this;
        }
        
        public void close() throws IOException {
            
            if (isNio) {
                nio.close();
            }
            else {
                aio.close();
            }
        }
    }
}
